package com.alatus.djl.model;

import ai.djl.Model;
import ai.djl.nn.Block;
import ai.djl.nn.SequentialBlock;
import ai.djl.nn.core.Linear;
import ai.djl.nn.Activation;

/*
* @author: alatuslee@qq.com
* 自编码器模型构建
* */
public class AutoEncoderModel {

    public static Model getAutoEncoderModel() {
        Model model = Model.newInstance("hardness-autoencoder");
        model.setBlock(createAutoEncoderBlock());
        return model;
    }

    private static Block createAutoEncoderBlock() {
        SequentialBlock block = new SequentialBlock();

        // 编码器部分
        block.add(Linear.builder().setUnits(128).build());
        block.add(Activation::relu);
        block.add(Linear.builder().setUnits(64).build());
        block.add(Activation::relu);
        block.add(Linear.builder().setUnits(32).build());  // 瓶颈层 - 特征提取
        block.add(Activation::relu);

        // 解码器部分
        block.add(Linear.builder().setUnits(64).build());
        block.add(Activation::relu);
        block.add(Linear.builder().setUnits(128).build());
        block.add(Activation::relu);
        block.add(Linear.builder().setUnits(252).build()); // 重构到原始维度

        return block;
    }

    // 仅编码器部分，用于特征提取
    public static Block getEncoderBlock() {
        SequentialBlock encoder = new SequentialBlock();

        encoder.add(Linear.builder().setUnits(128).build());
        encoder.add(Activation::relu);
        encoder.add(Linear.builder().setUnits(64).build());
        encoder.add(Activation::relu);
        encoder.add(Linear.builder().setUnits(32).build()); // 特征维度
        encoder.add(Activation::relu);

        return encoder;
    }
}